Mooncake: Kimi’s KVCache-centric Architecture for LLM Serving](https://arxiv.org/abs/2407.00079)

月之暗面kimi团队在2024的一篇工作

Mooncake: Trading More Storage for Less Computation — A KVCache-centric Architecture for Serving LLM Chatbot | USENIX

FAST 2025 Best Paper


Abstract

  • 以KV cache为核心
  • 采用了prefill和decode的分解架构
  • 利用了集群中未被使用的CPU、DRAM和SSD等资源
  • 在最大化吞吐量的同时,还平衡了服务级别目标
  • 面临的是高度超载的环境,提出了基于预测的早期拒绝政策,
  1. Model as a Service (MaaS)的最大目标是 maximize overall effective throughput和尽量约束varying levels of SLOs,比如 the time to first token (TTFT)和 the time between tokens (TBT)。
  2. prefill和decoding计算特性不一样,前者是compute bound,后者是memory bound,曾有相关工作提出了分离架构。

以KV Cache为中心的架构

image-20250305155814860

分为三部分:

  • Prefill Pool
    • max Cache Reuse
  • Mooncake Store
  • Decoding Pool
    • max Throughput

Workfolw

image-20250305160034874

  1. KVCache Reuse

    • 尽可能多地reuse KV Cache:一边计算A一边传输B,可以理解为减轻计算负载;同时还能减低TTFT时间
    • 平衡各节点负载
    • 满足TTFT
  2. Incremental Prefill

    • 以块的形式计算和传输KV Cache,减少KV Cache传输造成的阻塞时间
  3. KVCache Transfer

    • 和Prefill异步,从而实现overlap
  4. Decoding

    • 根据instance负载选择合适的decoding instance

    • 当KV Cache传输到Decoding Instance完成时,加入下一iteration的计算

Mooncake Store

【3.2.1】KV Cache管理:存储的元数据,比如说前缀信息HashMap

【3.2.2】API接口介绍

【3.2.3】设计传输引擎

目标:

  1. 有效使用多个RDMA NIC
  2. 抽象出来的API
  3. 支持错误处理

NCCL不支持动态调整和DRAM to DRAM的路径

1. Topology-aware path selection

image-20250305161612072

  1. 基于拓扑的链接

    1. 广播拓扑
    2. 建立链接
  2. 分块传输(16KB)

    2. Endpoint Managerment

维护了一个连接的Endpoint Pool,按需生成Endpoint,Pool限制最大数量,采用了SIEVE算法。

3. Failure Handing

自动识别损坏的路径并更新拓扑

Mooncake‘s Prefill Pool

Chunked pipeline parallelism

Cache-aware的全局调度策略

  1. 基于请求长度、prefix cache命中率和local queuing time,计算各个新请求在不同Instance上的TTFT,分配给TTFT最小的
  2. TTFT最小的节点不一定是持有最大前缀匹配的instance,该instance可以主动地向持有其最大前缀匹配的节点获取KV Cache数据。
  3. 如果本地的前缀匹配token数乘以一个阈值大于所有节点的最大前缀匹配token数,直接本地计算。

image-20250319143705620

跑了16个8*A100 nodes的实验证明,这种方式TTFT更优。

实验结果

image-20250319144118803

在不同输入和更长的输出情况下(Chatbot),vLLM TBT违反率很高。

image-20250319144457392

在较短的输入和输出时,vLLM多节点前缀利用率不如Mooncake。

image-20250319144723952

在混合各种dataset的实验,Mooncake表现依旧更好。

Prefill时间实验。

利用global cache后的前缀缓存命中率实验:

image-20250319145100988

cache-aware机制下缓存块副本的数量:比较稳定

image-20250319145616769

通信实验:

image-20250319145744373

100Gbps是Mooncake这种Cache机制下的底线

image-20250319150023523

前缀缓存降低了很多时间,虽然其他部分时间增长,但前缀缓存还是有收益

image-20250319145917713

PD实验:1比1 其实比较稳定,可以在不同时间段动态切换一下

image-20250319150422244

results matching ""

    No results matching ""